#!/usr/bin/env python3
import csv
import math
from collections import defaultdict

in_path = '/home/ubantu00/zk_demo/build/results.csv'
out_txt = '/home/ubantu00/zk_demo/build/results_summary.txt'
out_csv = '/home/ubantu00/zk_demo/build/results_summary.csv'

# data structure: key -> list of rows
# row: dict with keys from header
groups = defaultdict(list)
with open(in_path, newline='') as f:
    reader = csv.DictReader(f)
    for r in reader:
        key = (r['scenario'], r['type'])
        # convert numeric fields
        try:
            r['prove_ms'] = float(r['prove_ms'])
        except:
            r['prove_ms'] = float('nan')
        try:
            r['verify_ms'] = float(r['verify_ms'])
        except:
            r['verify_ms'] = float('nan')
        def safe_int(x):
            try:
                if x is None:
                    return 0
                s = str(x).strip().replace(',','')
                return int(s) if s != '' else 0
            except:
                return 0
        r['proof_bytes'] = safe_int(r.get('proof_bytes', 0))
        r['vk_bytes'] = safe_int(r.get('vk_bytes', 0))
        r['verifier_recv_bytes'] = safe_int(r.get('verifier_recv_bytes', 0))
        try:
            r['comm_bytes'] = int(r['comm_bytes'])
        except:
            r['comm_bytes'] = 0
        groups[key].append(r)

# compute stats
summary_rows = []
for (scenario, typ), rows in sorted(groups.items()):
    n = len(rows)
    prove_vals = [r['prove_ms'] for r in rows if not math.isnan(r['prove_ms'])]
    verify_vals = [r['verify_ms'] for r in rows if not math.isnan(r['verify_ms'])]
    proof_bytes_vals = [r['proof_bytes'] for r in rows]
    vk_bytes_vals = [r['vk_bytes'] for r in rows]
    verifier_recv_vals = [r['verifier_recv_bytes'] for r in rows]
    total_comm = sum(r['comm_bytes'] for r in rows)
    total_verifier_recv = sum(r['verifier_recv_bytes'] for r in rows)

    def mean(xs):
        return sum(xs)/len(xs) if xs else float('nan')
    def std(xs):
        if len(xs) < 2:
            return 0.0
        m = mean(xs)
        return math.sqrt(sum((x-m)**2 for x in xs)/(len(xs)-1))

    row = {
        'scenario': scenario,
        'type': typ,
        'n': n,
        'prove_mean_ms': mean(prove_vals),
        'prove_std_ms': std(prove_vals),
        'verify_mean_ms': mean(verify_vals),
        'verify_std_ms': std(verify_vals),
        'proof_mean_bytes': mean(proof_bytes_vals),
        'vk_mean_bytes': mean(vk_bytes_vals),
        'verifier_recv_total_bytes': total_verifier_recv,
        'total_comm_bytes': total_comm,
    }
    summary_rows.append(row)

# write text and csv
with open(out_txt, 'w') as f:
    for r in summary_rows:
        f.write(f"Scenario: {r['scenario']} ({r['type']})\n")
        f.write(f"  samples: {r['n']}\n")
        f.write(f"  prove mean(ms): {r['prove_mean_ms']:.6f}, std(ms): {r['prove_std_ms']:.6f}\n")
        f.write(f"  verify mean(ms): {r['verify_mean_ms']:.6f}, std(ms): {r['verify_std_ms']:.6f}\n")
        f.write(f"  proof mean(bytes): {r['proof_mean_bytes']:.1f}\n")
        f.write(f"  total comm bytes: {r['total_comm_bytes']}\n\n")

with open(out_csv, 'w', newline='') as f:
    writer = csv.writer(f)
    writer.writerow(['scenario','type','n','prove_mean_ms','prove_std_ms','verify_mean_ms','verify_std_ms','proof_mean_bytes','vk_mean_bytes','verifier_recv_total_bytes','total_comm_bytes'])
    for r in summary_rows:
        writer.writerow([r['scenario'], r['type'], r['n'],
                         f"{r['prove_mean_ms']:.6f}", f"{r['prove_std_ms']:.6f}",
                         f"{r['verify_mean_ms']:.6f}", f"{r['verify_std_ms']:.6f}",
                         f"{r['proof_mean_bytes']:.1f}", f"{r['vk_mean_bytes']:.1f}", r['verifier_recv_total_bytes'], r['total_comm_bytes']])

print('Wrote', out_txt, 'and', out_csv)
